from typing import List

from elasticsearch import AsyncElasticsearch
from elasticsearch.helpers import async_bulk, async_scan
from loguru import logger
from pydantic import parse_obj_as

from config import settings
from .basic import es


async def connect_to_es():
    es.clients = dict()
    for name, es_conf in settings.es.items():
        logger.info(f"init elasticsearch client {name}")
        es.clients[name.lower()] = AsyncElasticsearch(f"{es_conf['host']}:{es_conf['port']}", http_auth=(es_conf['username'], es_conf['password']), timeout=1000)


async def close_es_connect():
    for name, client in es.clients.items():
        logger.info(f"close elasticsearch client {name}")
        await client.close()


def add_highlight2source(source, hit_item):
    # TODO 待优化
    highlight = hit_item.get('highlight', {})
    for k, v in highlight.items():
        highlight[k] = v[0]
    if highlight:
        source['highlight'] = highlight


async def search_models(index_const, model_type, query, **params):
    client, index = es.get_connect_info(index_const)
    result = await client.search(
        index=index,
        body=query,
        **params,
    )
    docs = []
    total = result["hits"]['total']['value']
    for item in result["hits"]['hits']:
        source = item['_source']
        source['id'] = item['_id']
        source['score'] = item['_score']
        add_highlight2source(source, item)
        docs.append(source)
    models = parse_obj_as(List[model_type], docs)
    aggs = result.get('aggregations')
    return total, models, aggs


async def search_docs(index_const, query, **params):
    client, index = es.get_connect_info(index_const)
    resp = await client.search(
        index=index,
        body=query,
        **params,
    )
    return resp


async def scan_docs(index_const, query, **params):
    """
    用于获取超过10000个文档数据
    """
    client, index = es.get_connect_info(index_const)
    doc_generator = async_scan(client, query=query, index=index, **params)
    return doc_generator


async def update_doc(index_const, doc_id, body, **params):
    client, index = es.get_connect_info(index_const)
    resp = await client.update(index, doc_id, body, **params)
    return resp


async def create_doc(index_const, doc_id, body):
    client, index = es.get_connect_info(index_const)
    resp = await client.index(index=index, doc_type="_doc", id=doc_id, body=body)
    return resp


async def bulk_update_docs(index_const, docs, stats_only=False, raise_on_error=True):
    client, index = es.get_connect_info(index_const)
    bulk_actions = list()
    for doc in docs:
        bulk_actions.append({
            '_op_type': 'update',
            '_index': index,
            '_id': doc.pop('id'),
            'doc': doc,
        })
    return await async_bulk(client, bulk_actions, stats_only=stats_only, raise_on_error=raise_on_error)


async def bulk_upsert_docs(index_const, docs, stats_only=False):
    client, index = es.get_connect_info(index_const)
    bulk_actions = list()
    for doc in docs:
        bulk_actions.append({
            '_op_type': 'update',
            '_index': index,
            '_id': doc.pop('id'),
            'doc': doc,
            'doc_as_upsert': True,
        })
    return await async_bulk(client, bulk_actions, stats_only=stats_only)


def model2doc(model, **kwargs):
    exclude_unset = kwargs.pop('exclude_unset', True)
    by_alias = kwargs.pop('by_alias', True)

    parsed = model.dict(
        exclude_unset=exclude_unset,
        by_alias=by_alias,
        **kwargs,
    )
    return mid, parsed


async def index_one_model(index_const, model):
    """
    索引一个model
    """
    mid, doc = model2doc(model, exclude_unset=True)
    client, index = es.get_connect_info(index_const)
    result = await client.index(index, id=mid, body=doc)
    return result


async def index_many_models(index_const, models, stats_only=False):
    """
    批量索引model
    """
    client, index = es.get_connect_info(index_const)
    bulk_actions = list()
    for m in models:
        mid, doc = model2doc(m, exclude_unset=True)
        doc['_index'] = index
        doc['_id'] = mid
        bulk_actions.append(doc)
    return await async_bulk(client, bulk_actions, stats_only=stats_only)


async def update_many_models(index_const, models, stats_only=False, raise_on_error=True):
    client, index = es.get_connect_info(index_const)
    bulk_actions = list()
    for m in models:
        mid, doc_body = model2doc(m, exclude_unset=True)
        bulk_actions.append({
            '_op_type': 'update',
            '_index': index,
            '_id': mid,
            'doc': doc_body
        })
    return await async_bulk(client, bulk_actions, stats_only=stats_only, raise_on_error=raise_on_error)


async def upsert_models(index_const, models, stats_only=False):
    client, index = es.get_connect_info(index_const)
    bulk_actions = list()
    for m in models:
        mid, doc_body = model2doc(m, exclude_unset=True)
        bulk_actions.append({
            '_op_type': 'update',
            '_index': index,
            '_id': mid,
            'doc': doc_body,
            'doc_as_upsert': True,
        })
    return await async_bulk(client, bulk_actions, stats_only=stats_only)


async def delete_by_query(index_const, query):
    client, index = es.get_connect_info(index_const)
    result = await client.delete_by_query(index, body=query)
    return result


async def count_documents(index_const, query):
    client, index = es.get_connect_info(index_const)
    result = await client.count(index=index, body=query)
    return result


async def create_index(client_name, index_name, settings):
    client = es[client_name]
    result = await client.indices.create(index=index_name, body=settings)
    return result


async def delete_index(client_name, index_name):
    client = es[client_name]
    result = await client.indices.delete(index=index_name)
    return result
